import cv2
import numpy as np
from numba import cuda,njit
from numba import int64 as nint64
from numba.cuda.random import create_xoroshiro128p_states,xoroshiro128p_uniform_float32
import time
import random
import numba
# 常量
DEFAULT_NUM_SAMPLES = 20
DEFAULT_THREAD_PER_BLOCK = 32
DEFAULT_RADIUS = 15
DEFAULT_MIN_MATCHES = 2
DEFAULT_RANDOM_SAMPLE = 16
DEFAULT_MIS_COUNT = 50
# LCG 参数
A = 1664525
C = 1013904223
M = 2**32

h_x_shift = np.array([-1, 0, 1, -1, 0, 1, -1, 0, 1], dtype=np.int32)
h_y_shift = np.array([-1, 0, 1, -1, 0, 1, -1, 0, 1], dtype=np.int32)
d_x_shift = cuda.to_device(h_x_shift)
d_y_shift = cuda.to_device(h_y_shift)


COLORS = [
        [56, 0, 255],
        [226, 255, 0],
        [0, 94, 255],
        [0, 37, 255],
        [0, 255, 94],
        [255, 226, 0],
        [0, 18, 255],
        [255, 151, 0],
        [170, 0, 255],
        [0, 255, 56],
        [255, 0, 75],
        [0, 75, 255],
        [0, 255, 169],
        [255, 0, 207],
        [75, 255, 0],
        [207, 0, 255],
        [37, 0, 255],
        [0, 207, 255],
        [94, 0, 255],
        [0, 255, 113],
        [255, 18, 0],
        [255, 0, 56],
        [18, 0, 255],
        [0, 255, 226],
        [170, 255, 0],
        [255, 0, 245],
        [151, 255, 0],
        [132, 255, 0],
        [75, 0, 255],
        [151, 0, 255],
        [0, 151, 255],
        [132, 0, 255],
        [0, 255, 245],
        [255, 132, 0],
        [226, 0, 255],
        [255, 37, 0],
        [207, 255, 0],
        [0, 255, 207],
        [94, 255, 0],
        [0, 226, 255],
        [56, 255, 0],
        [255, 94, 0],
        [255, 113, 0],
        [0, 132, 255],
        [255, 0, 132],
        [255, 170, 0],
        [255, 0, 188],
        [113, 255, 0],
        [245, 0, 255],
        [113, 0, 255],
        [255, 188, 0],
        [0, 113, 255],
        [255, 0, 0],
        [0, 56, 255],
        [255, 0, 113],
        [0, 255, 188],
        [255, 0, 94],
        [255, 0, 18],
        [18, 255, 0],
        [0, 255, 132],
        [0, 188, 255],
        [0, 245, 255],
        [0, 169, 255],
        [37, 255, 0],
        [255, 0, 151],
        [188, 0, 255],
        [0, 255, 37],
        [0, 255, 0],
        [255, 0, 170],
        [255, 0, 37],
        [255, 75, 0],
        [0, 0, 255],
        [255, 207, 0],
        [255, 0, 226],
        [255, 245, 0],
        [188, 255, 0],
        [0, 255, 18],
        [0, 255, 75],
        [0, 255, 151],
        [255, 56, 0],
        [245, 255, 0],
    ]

def get_crop_info(zone):
        (x, y), (x1, y1) = np.min(zone, axis=0).flatten(), np.max(zone, axis=0).flatten()
        y1 = 1080 if y1 > 1080 else y1
        x1 = 1920 if y1 > 1920 else x1
        w, h = x1 - x, y1 - y
        return [int(x), int(y), int(w), int(h)]

def slice_numpy_img(img,box):
        box = (int(i) for i in box)
        x,y,w,h = box
        return img[y:y+h,x:x+w]

def create_rng_states(h, w):
    return np.random.randint(0, 2**32, size=(h, w), dtype=np.uint32)


@cuda.jit
def get_vibe_background(img_init, img_bg, h_bg, w_bg, c_bg,d_roi_mat,rng_states,d_x_shift,d_y_shift):
    idx, idy, idz = cuda.grid(3)
    if idx < h_bg and idy < w_bg and idz < c_bg :
        if d_roi_mat[idx,idy] ==1:

            # 从状态生成一个0到1之间的随机数
            random_float1 = xoroshiro128p_uniform_float32(rng_states, idx*w_bg+idy)
            # 将其缩放到所需范围
            random_num1 = int(0 + random_float1 * (9 - 0))

             # 从状态生成一个0到1之间的随机数
            random_float2 = xoroshiro128p_uniform_float32(rng_states, idx*w_bg+idy)
            # 将其缩放到所需范围
            random_num2 = int(0 + random_float2 * (9 - 0))

            # 将随机数限制在 0 到 10 之间
            random_shift_x = d_x_shift[random_num1]  # 0 到 9
            random_shift_y = d_y_shift[random_num2]  # 0 到 9

            # print("000 -1-1",random_shift_x,random_shift_y)
            # if 0 <= idx + random_shift_x < h_bg and 0 <= idy + random_shift_y < w_bg:
            random_offset_x = idx + random_shift_x
            if idx + random_shift_x<0:
                random_offset_x = 0
            elif idx + random_shift_x>=h_bg:
                random_offset_x = h_bg-1
            random_offset_y = idy + random_shift_y
            if idy + random_shift_y<0:
                random_offset_y = 0
            elif idy + random_shift_y>=w_bg:
                random_offset_y = w_bg-1
            # print("000 -1-1",random_offset_x,random_offset_y)
                
            img_bg[idx, idy, idz] = img_init[random_offset_x, random_offset_y]



@cuda.jit
def cal_match_count(img_cur, img_bg, img_count, threshold_distance,h, w, c, d_roi_mat):
    idx, idy, idz = cuda.grid(3)
    # 此处的c为样本库通道总数
    if idx < h and idy < w and idz<c :
        # 防区掩码过滤
        if d_roi_mat[idx,idy] ==1:
            # 距离小于阈值,匹配度+1
            if abs(img_bg[idx, idy, idz]-img_cur[idx, idy]) < threshold_distance:
                img_count[idx,idy] += 1
# @cuda.jit
# def cal_match_count(img_cur, img_bg, img_count, threshold_distance,h, w, c, d_roi_mat):
#     idx, idy, idz = cuda.grid(3)
#     # 此处的c为样本库通道总数
#     if idx < h and idy < w and idz<c :
#         # 防区掩码过滤
#         if d_roi_mat[idx,idy] ==1:
#             # 距离小于阈值,匹配度+1
#             match_count = 0
#             for idz in range(c):
#                 if abs(img_bg[idx, idy, idz]-img_cur[idx, idy]) < threshold_distance:
#                     match_count += 1
#             img_count[idx,idy] += match_count



@cuda.jit
def get_diff(img_cur, img_bg, img_count,target_hit_count, d_mask_img, h, w, d_roi_mat,rng_states1,min_match,ramdom_sample,channel_num,maxhitcount,d_x_shift,d_y_shift):
    idx, idy = cuda.grid(2)
    if idx < h and idy < w:
        if d_roi_mat[idx,idy] ==1:
            # for m in range(channel_num):
            #     if abs(img_bg[idx, idy, m]-img_cur[idx, idy]) < 30:
            #         img_count[idx,idy] += 1
            #     if img_count[idx,idy] == min_match:
            #         break
            if img_count[idx,idy] >= min_match:
                d_mask_img[idx, idy] = 0
                target_hit_count[idx,idy] = 0
                # random_num = random_num1 % (ramdom_sample+1)
                # 从状态生成一个0到1之间的随机数
                random_float0 = xoroshiro128p_uniform_float32(rng_states1, idx*w+idy)
                # 将其缩放到所需范围
                random_num = int(0 + random_float0 * (ramdom_sample - 0))
                
                # print("222  0-16",random_num)
                # print("random_num",random_num)

                if random_num==0:

                    # 从状态生成一个0到1之间的随机数
                    random_float1 = xoroshiro128p_uniform_float32(rng_states1, idx*w+idy)
                    # 将其缩放到所需范围
                    random_num1 = int(0 + random_float1 * (9 - 0))

                    # 从状态生成一个0到1之间的随机数
                    random_float2 = xoroshiro128p_uniform_float32(rng_states1, idx*w+idy)
                    # 将其缩放到所需范围
                    random_num2 = int(0 + random_float2 * (9 - 0))

                    random_shift_x = d_x_shift[random_num1]  # 0 到 9
                    random_shift_y = d_y_shift[random_num2]  # 0 到 9
                    # print("333 -1-1",random_shift_x,random_shift_y)
                    random_offset_x = idx + random_shift_x
                    if idx + random_shift_x<0:
                        random_offset_x = 0
                    elif idx + random_shift_x>=h:
                        random_offset_x = h-1
                    random_offset_y = idy + random_shift_y
                    if idy + random_shift_y<0:
                        random_offset_y = 0
                    elif idy + random_shift_y>=w:
                        random_offset_y = w-1
                    random_float3 = xoroshiro128p_uniform_float32(rng_states1, idx*w+idy)
                    # 将其缩放到所需范围
                    random_c3 = int(0 + random_float3 * (channel_num - 0))
                    # print("444 0-20 ",random_c1)

                    img_bg[random_offset_x,random_offset_y,random_c3] = img_cur[idx,idy]
                    
            else:
                d_mask_img[idx,idy] =255
                target_hit_count[idx,idy] += 1
                if target_hit_count[idx,idy]>=maxhitcount:
                    # random_c2 = random_num2 % (channel_num+1)
                    random_float4 = xoroshiro128p_uniform_float32(rng_states1, idx*w+idy)
                    # 将其缩放到所需范围
                    random_c4 = int(0 + random_float4 * (channel_num - 0))
                    # print("555 0-20 ",random_shift_x,random_shift_y)
                    img_bg[idx,idy,random_c4] = img_cur[idx,idy]
                    # d_mask_img[idx,idy] =0
                    # target_hit_count[idx,idy]  = 0


import numpy as np

def apply_mask_to_image(image, mask):
    """
    将mask应用到图像上，mask区域显示为黄色。
    
    参数:
    image: 原始彩色图像 (H, W, C)
    mask: 二值化的mask (H, W)，mask区域为255，非mask区域为0

    返回:
    masked_image: 应用mask后的彩色图像 (H, W, C)
    """
    # 确保输入图像为彩色图像
    if len(image.shape) == 2:  # 如果是灰度图像
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)

    # 创建一个与输入图像相同形状的全黑图像
    masked_image = np.zeros_like(image)

    # 将mask区域设为黄色 (BGR格式，黄色为(0, 255, 255))
    clolor = random.choice(COLORS)
    masked_image[mask == 255] = clolor # Yellow in BGR

    # 将原始图像的非mask区域保留
    masked_image[mask == 0] = image[mask == 0]

    return masked_image



RTSP = "2.mp4"
# zone1 = "860,249,1156,650,1292,588,959,249,-1"
zone1 = "0,0,1920,0,1920,1080,0,1920,-1"

# RTSP = "11.avi"
# zone1 = "580,1070,851,420,901,420,947,1068,-1"
zone = np.array(zone1.split(',-1')[0].split(','), dtype=np.int32).reshape((-1, 1, 2))
offset = get_crop_info(zone)
roi_mat = np.zeros((offset[3], offset[2]), dtype=np.uint8)
new_boundary = zone - np.array((offset[0], offset[1]))
cv2.fillPoly(roi_mat, [new_boundary], 1)
d_roi_mat = cuda.to_device(roi_mat)
             
        
vs = cv2.VideoCapture(RTSP)
ret,frame = vs.read()
# h, w,c = frame.shape
frame = cv2.cvtColor(frame,cv2.COLOR_BGR2GRAY)

frame = slice_numpy_img(frame,offset)*roi_mat
h, w = frame.shape
d_img_bg = cuda.to_device(frame.astype(np.float32))

bg = np.zeros((h, w, DEFAULT_NUM_SAMPLES), dtype=np.float32)
d_bg = cuda.to_device(bg)

# 定义网格和块的大小
block_per_grid_h = (h + DEFAULT_THREAD_PER_BLOCK - 1) // DEFAULT_THREAD_PER_BLOCK
block_per_grid_w = (w + DEFAULT_THREAD_PER_BLOCK - 1) // DEFAULT_THREAD_PER_BLOCK

# 创建网格
blocks_per_grid_bg = (block_per_grid_h, block_per_grid_w, DEFAULT_NUM_SAMPLES)
blocks_per_grid_diff = (block_per_grid_h, block_per_grid_w)

frame_id = 0

from collections import deque
result_quee = deque(maxlen=10)

img_match_count = np.zeros((h, w), dtype=np.int32)
d_img_match_count = cuda.to_device(img_match_count)
# rng_states = create_rng_states(h, w)
rng_states = create_xoroshiro128p_states(h*w, seed=frame_id)  # 使用1234作为随机种子
get_vibe_background[blocks_per_grid_bg, (DEFAULT_THREAD_PER_BLOCK, DEFAULT_THREAD_PER_BLOCK, 1)](d_img_bg, d_bg, h, w, DEFAULT_NUM_SAMPLES,d_roi_mat,rng_states,d_x_shift,d_y_shift)

# 被认为是目标的计数次数
target_hit_count = np.zeros((h, w), dtype=np.uint8)
d_target_hit_count = cuda.to_device(target_hit_count)
rng_states = create_xoroshiro128p_states(h*w, seed=frame_id)  # 使用1234作为随机种子
rng_states1 = create_xoroshiro128p_states(h*w, seed=frame_id+1)  # 使用1234作为随机种子

# 二值化前景背景判别图
mask_img0 = np.zeros((h, w), dtype=np.uint8)
d_mask_img = cuda.to_device(mask_img0)

out = cv2.VideoWriter('cuda_vibe.avi',cv2.VideoWriter_fourcc('M','J','P','G'), 25.0, (1920,1080))
if __name__ == "__main__":
    while True:
        frame_id +=1
        # if frame_id % 5 == 0:
        #     continue
        ret, frame = vs.read()
        if not ret:
            break
        t = time.time()
        # if frame_id % 100 == 0:
        #     rng_states = create_xoroshiro128p_states(h*w, seed=frame_id)  # 使用1234作为随机种子
        #     rng_states1 = create_xoroshiro128p_states(h*w, seed=frame_id+1)  # 使用1234作为随机种子
        
        # # 二值化前景背景判别图
        # mask_img0 = np.zeros((h, w), dtype=np.uint8)
        # d_mask_img = cuda.to_device(mask_img0)

        # 被认为是背景的计数次数
        img_match_count = np.zeros((h, w), dtype=np.uint8)
        d_img_match_count = cuda.to_device(img_match_count)

        # 灰度图计算模式: 裁剪防区 && 防区外掩码
        frame_crop = slice_numpy_img(frame,offset)
        frame_gray = cv2.cvtColor(frame_crop,cv2.COLOR_BGR2GRAY)
        frame_mask = frame_gray*roi_mat
        # print("frame,d_bg",frame,d_bg)
        d_img_cur = cuda.to_device(frame_mask.astype(np.float32))

        # 计算背景匹配数
        t1 = time.time()
        cal_match_count[blocks_per_grid_bg, (DEFAULT_THREAD_PER_BLOCK, DEFAULT_THREAD_PER_BLOCK, 1)](d_img_cur, d_bg, d_img_match_count, DEFAULT_RADIUS,h, w, DEFAULT_NUM_SAMPLES, d_roi_mat)
        # cuda.synchronize()  # 确保所有CUDA操作完成
        
        t2 = time.time()
        print(f"cal match count: {t2 - t1:.4f} s")
        get_diff[blocks_per_grid_diff, (DEFAULT_THREAD_PER_BLOCK, DEFAULT_THREAD_PER_BLOCK)](d_img_cur, d_bg, d_img_match_count,d_target_hit_count, d_mask_img, h, w,d_roi_mat,rng_states1,
                                                                                            DEFAULT_MIN_MATCHES,DEFAULT_RANDOM_SAMPLE,DEFAULT_NUM_SAMPLES,DEFAULT_MIS_COUNT,d_x_shift,d_y_shift)
        cuda.synchronize()  # 确保所有CUDA操作完成
        # d_img_match_count.copy_to_host(img_match_count)
        # print(img_match_count)

        t3 = time.time()
        print(f"get diff: {t3 - t2:.4f} s")
        mask_sult = d_mask_img.copy_to_host()
        t4 = time.time()
        print(f"device2host: {t4 - t3:.4f} s")

        # result_quee.append(mask_sult)
        # 将deque转换为NumPy数组（最后一个维度是图像通道）
        # images_array = np.stack(result_quee, axis=-1)

        # # 执行与运算
        # result_image = np.bitwise_and.reduce(images_array, axis=-1)
        t5 = time.time()
        print(f"total cast {t5-t1} FPS {1/(t5-t1)}")
        # if cv2.waitKey(10) & 0xff == ord("q"):
        #     out.release()
        #     break
            
        # mask_sult = cv2.medianBlur(mask_sult,5)
        # mask_sult[mask_sult>0]=255
        # # mask_sult = cv2.blur(mask_sult,(5,5))
        
        # global_mask = np.zeros((frame.shape[0],frame.shape[1])).astype(np.uint8)
        # global_mask[offset[1]:offset[1]+offset[3],offset[0]:offset[0]+offset[2]] = mask_sult
        # result_image0 = apply_mask_to_image(frame,global_mask)
        # # cv2.waitKey(10)
        # cv2.namedWindow("VIBE_CUDA",cv2.WINDOW_NORMAL)
        # cv2.putText(result_image0, f"FPS: {int(1/(t5-t))}", (10, 100), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 255, 0), 5)
        # cv2.imshow("VIBE_CUDA",result_image0)
        # out.write(result_image0.astype(np.uint8))
        # print(result_image0.shape)
        # 选择GPU设备
    out.release()
    numba.cuda.select_device(0)

    # 关闭GPU设备，释放所有资源
    numba.cuda.close()
    
